'''
Created on Feb 2, 2011

:author: Ruyman Reyes (rreyes@ull.es)



'''

from yacf.Frontend.Shortcuts import getCurrentLanguageAst
c_ast = getCurrentLanguageAst()


from yacf.Tools.Tree import getFileNode
from yacf.Tools.Declarations import decl_of_id 

from yacf.Backends.C.Exceptions import NonOmpCompilant
from yacf.Backends.C99.Visitors.LlcFilters import AccScopeReverse
from yacf.Backends.C99.Visitors.OmpFilters import OmpScopeReverse
from yacf.Frontend.IRTypes import DeclList
from yacf.Backends.Common.Visitors.GenericVisitors import UndefIDFilter
from yacf.Frontend.SymbolTable import getSymbolTableFromNode

from yacf.Tools.Tree import are_equivalent

class ParametrizeLoopTool(object):
    """ Extract information from a canonical Loop.
            Canonical Loops are defined on the OpenMP 3.0, p 38
    """
    def __init__(self,):
        self._parameters = {}
        self.equals = 0

    def symbolic_max(self, node_ast):
        """ Recursively estimate the maximum possible value of an expression

        Note: Assuming the local maximum is the global one

		:param node_ast: expression node 
			from where extract information

		:return: Symbolic maximum value of the expression.

	"""
        if isinstance(node_ast, c_ast.BinaryOp):
            if node_ast.op in ['<', '<='] :
                if node_ast.op == '<=':
                    self.equals = self.equals + 1
                return self.symbolic_max(node_ast.right)
            elif node_ast.op in ['>', '>=']:
                if node_ast.op == '>=':
                    self.equals = self.equals + 1
                return self.symbolic_max(node_ast.left)
            elif node_ast.op in ['-', '+', '*', '/']:
                # The result of this operation is treated as a whole element
                return node_ast
            elif node_ast.op == '==':
                # Both are maximums, this is a situation with two local maximums
                raise NotImplemented
            else:
                raise NotImplemented
        elif isinstance(node_ast, c_ast.UnaryOp):
            if node_ast.op == '*':
                return node_ast;
            else:
                raise NotImplemented
        elif isinstance(node_ast, c_ast.ID):
            return node_ast
        elif isinstance(node_ast, c_ast.Constant):
            """ This constant should be the maximum value """
            return node_ast
        elif isinstance(node_ast, c_ast.TernaryOp):
            return node_ast
        else:
            import pdb
            pdb.set_trace()
            raise NotImplemented
    

    
    def get_loop_variable(self, loop_ast):
	""" Extracts the iterating variable (loop variable) of
		the given loop statement. 
        	Loop must be on canonical form (as in page 39, line 42 of OpenMP 3.0
		standard)

		Note: For now, we assume that the loop is a for loop

		:param loop_ast: Loop statement node
		:return: Node with the ID of the loop variable
	"""
        if isinstance(loop_ast.init, c_ast.Assignment):
            return loop_ast.init.lvalue
        elif isinstance(loop_ast.init, c_ast.DeclList):
            # Assuming only one declaration
            if len(loop_ast.init.decls) == 1:
                # This is a declaration
                decl = loop_ast.init.decls[0]
                return c_ast.ID(name = decl.name, 
				sequence = decl.sequence, 
				depth = decl.depth,
				parent = loop_ast.init)
            else:
                print "More than one declaration, not implemented"
        elif hasattr(loop_ast.init, 'left'):
            return loop_ast.init.left
        raise NonOmpCompilant('Loop not on canonical form, see Page 39, line 42 of OpenMP 3.0 Standard')
    
    def get_init_value(self, init):
	""" Extracts the initial value of the loop variable of
		the given loop statement. 
        	Loop must be on canonical form (as in page 39, line 42 of OpenMP 3.1
		standard)

		Note: For now, we assume that the loop is a for loop

		:param loop_ast: Loop statement node
		:return: Node with the ID of the loop variable
	"""
        if isinstance(init, c_ast.Assignment):
            return init.rvalue
        elif hasattr(init, 'right'):
            return init.right
        elif isinstance(init, c_ast.DeclList):
            # Assuming only one declaration
            if len(init.decls) == 1:
                # This is a declaration
                decl = init.decls[0]
                return decl.init
            else:
                print "More than one declaration, not implemented"
        raise NonOmpCompilant('Loop not on canonical form, see Page 42 of OpenMP 3.0 Standard')
        
    def get_stride(self, next_ast, loop_var):
	""" Extracts the stride of the given statement using loop_var
		the given loop statement. 
        	Loop must be on canonical form (as in page 39, line 42 of OpenMP 3.0
		standard)

		Note: For now, we assume that the loop is a for loop

		:param loop_ast: Loop statement node
		:return: Node with the ID of the loop variable
	"""
        if  isinstance(next_ast, c_ast.UnaryOp):
            if next_ast.op == 'p++' or next_ast.op == '++':
                return c_ast.Constant(type = 'int', value = "1", )
            elif next_ast.op == 'p--' or next_ast.op == '--':
                return c_ast.Constant(type = 'int', value = "-1")
        elif isinstance(next_ast, c_ast.BinaryOp):
            if next_ast.op == '+' and are_equivalent(next_ast.left, loop_var):
                return next_ast.right 
            elif next_ast.op == '+' and are_equivalent(next_ast.right,loop_var):
                return next_ast.left
            elif next_ast.op == '-' and are_equivalent(next_ast.left, loop_var):
                return next_ast.right 
        elif isinstance(next_ast, c_ast.Assignment):
            if next_ast.op == '+=' and are_equivalent(next_ast.lvalue, loop_var):
                return next_ast.rvalue
            elif next_ast.op == '-=' and are_equivalent(next_ast.lvalue, loop_var):
                return next_ast.rvalue
            elif next_ast.op == '=' and are_equivalent(next_ast.lvalue, loop_var):
                return self.get_stride(next_ast.rvalue, loop_var)
        # Implemented every case of the standard.
        import pdb
        pdb.set_trace()
        raise NonOmpCompilant('Loop not on canonical form, see Page 42 of OpenMP 3.0 Standard')
    

    def _compute_calculated_parameters(self):
        self.equals = 0
        self._parameters['last_iteration'] = self.symbolic_max(self._parameters['cond_node'],)
        self._parameters['init_iteration'] = self.get_init_value(self._parameters['init_node'])
        self._parameters['number_of_iterations'] = c_ast.BinaryOp(op='+', right = c_ast.BinaryOp(op='/', left=c_ast.BinaryOp(
                op = '-', 
                left = self._parameters['last_iteration'], 
                right = self._parameters['init_iteration']), 
            right = self._parameters['stride']),
            left = c_ast.Constant(type = 'integer', value = str(self.equals))
            )
        self._parameters['iterator'] = c_ast.UnaryOp(
                    op='p++', 
                    expr = self._parameters['loop_variable']) \
                if (self._parameters['stride'] == 1) \
                else c_ast.Assignment(op='=', 
                    lvalue=self._parameters['loop_variable'], 
                    rvalue=c_ast.BinaryOp(
                        op='+', 
                        left=self._parameters['loop_variable'], 
                        right=self._parameters['stride'])
            )

    def parametrize(self, for_loop):
        """ Extracts loop information from the for_loop statement.
	
       Following is a list of dictionary value/key pars returned by :meth:`parametrize`, as it defined on ver. 1.0.4. 
       Future versions may add additional parameters.

       * loop_variable: Contains the variable that this loop is iterating
       * init_node: The ast node with the for loop initialization
       * cond_node: The ast node with the condition of the loop
       * init_iteration: First value of the loop_variable at which this loop will start
       * last_iteration: Final value of the loop_variable at which this loop will stop
       * stride: Distance of the loop_var among two iterations
       * ndims : Number of dimensions of the loop. Used for nested loop (two nested loops will set ndim to 2). Currently fixed to 1.
       * number_of_iterations: Ast node containing a expression that will be solved on execution time. This expression compute the number of iterations that this loop will perform given the previously extracted parameters.

        The method expects a c_ast.For node.
        
		:param loop_ast: Loop statement node
		:return: Dictionary with loop parameters, as detailed above.
    	"""
        assert isinstance(for_loop, c_ast.For)
        
        
        self._parameters['loop_variable'] = self.get_loop_variable(for_loop)
        self._parameters['init_node'] = for_loop.init
        self._parameters['cond_node'] = for_loop.cond
        self._parameters['used_vars'] = self.extract_used_vars(for_loop)
        self._parameters['stride'] = self.get_stride(for_loop.next, self._parameters['loop_variable'])
        self._parameters['ndims'] = 1
        
        # Computed parameters
        self._compute_calculated_parameters()
                                        
        return self._parameters


    def extract_used_vars(self, for_loop_subtree):
        """ Return a list of variables used but not declared
            on the current subtree
        """
        nondecls = {}
        target_node = for_loop_subtree
        for elem in UndefIDFilter(container_node = target_node).dfs_iter(target_node):
            nondecls[elem.name] = elem
      
        symbols = []
      
        # Get the symbol table from any node (they share the same FileAST)
        if len(nondecls.values()) == 0:
            return symbols
        
        st = getSymbolTableFromNode(nondecls.values()[0])
        
        for node in nondecls.values():    
                symbols += [st.lookUp(node),]

        return symbols


    def apply(self, for_loop):
        return self.parametrize(for_loop)
    
    def update(self, new_parameters):
        """ Update the stored parameter array """
        self._parameters.update(new_parameters)
        # Keep consistency of the information
        self._compute_calculated_parameters()
    
    def buildLoop(self,):
        """ Write a loop from the current parameters """
        if not self._parameters:
            # This should raise an exception,
            # "Parameters have not been calculated"
            raise NotImplemented
        return c_ast.For(init = self._parameters['init_node'], 
                         cond = self._parameters['cond_node'], 
                         next = self._parameters['iterator'],
                         stmt = None,
                        )


    def __str__(self,):
        cad = ""
        for param in self._parameters:
            cad += str(param) + " : [" 
            if type(param) == type([]):
                for elem in self._parameters[param]:
                    cad += str(elem) + ", "
            else:
                cad += str(self._parameters[param])
            cad += "] \n"
        return cad


class OmpParametrize(ParametrizeLoopTool):
    """ Parametrize a loop using the information from Omp pragma.
    
   
    
    """
    
    @staticmethod
    def intersect(a, b):
	""" Support method to intersect to lists
		From an online python recipe.
	"""
        return list(set(a) & set(b))
    
    
    def  _get_scope_clause_dict(self, node,):
	""" Builds the OpenMP scope dictionary 
		:param node: Node from which extract information
		:return: Dictionary of clauses
	"""
        ompfilt = OmpScopeReverse()
        #------------------------------------------------------------ import pdb
        #------------------------------------------------------- pdb.set_trace()
        ompfilt.apply(node)
        return ompfilt.clauses
                
    
    def parametrize(self, for_loop,):
        fileAst = getFileNode(for_loop)
        basic_dir = super(OmpParametrize,self).parametrize(for_loop)
        clause_dict = self._get_scope_clause_dict(for_loop)
        if clause_dict.has_key('COLLAPSE'):
            basic_dir['ndims'] = clause_dict['COLLAPSE']
#            print " Collapse : " + basic_dir['ndims']
        else:
            basic_dir['ndims'] = 1
         
        # Not for now
        #  OmpParametrize.intersect(clause_dict['SHARED'],[]) 
        # Loop variable is always private
        basic_dir['local_vars'] = clause_dict['PRIVATE'] #  + [decl_of_id(basic_dir['loop_variable'], fileAst)]
        # TODO: For now, assume only one reduction clause with one operand
        basic_dir['reduction_vars'] = clause_dict['REDUCTION']
        # All variables by default
        # basic_dir['in_vars'] = clause_dict['COPY_IN'] # basic_dir['global_vars']
        # basic_dir['out_vars'] = clause_dict['COPY_OUT']
        basic_dir['in_vars'] = basic_dir['used_vars']
        basic_dir['out_vars'] = basic_dir['used_vars']
        basic_dir['global_vars'] = DeclList(clause_dict['SHARED']) & DeclList(basic_dir['in_vars']) & DeclList(basic_dir['out_vars']) - DeclList([getSymbolTableFromNode(fileAst).lookUp(basic_dir['loop_variable'])])
        return basic_dir

    def apply(self, for_loop,):
        return self.parametrize(for_loop)
    
    
